import cv2
import os
import numpy as np

# 文件夹路径
image_dir = "/media/champrin/Data/Desktop/rune dataset/extract/all2_yolo/set"  # 图片文件夹路径
output_dir = "/media/champrin/Data/Desktop/rune dataset/extract/all2_yolo/set/output"  # 输出图像文件夹路径
os.makedirs(output_dir, exist_ok=True)

# 绘图参数
drawing = False
last_point = None
mouse_x, mouse_y = -1, -1
circle_radius = 10

# 缩放参数
scale = 1.0
scale_step = 0.1
window_width, window_height = 800, 600
offset_x, offset_y = 0, 0  # 图像偏移量

# 撤回功能相关
history = []  # 存储历史图像状态
max_history = 10  # 最大历史记录数


def save_to_history(img):
    """保存当前图像状态到历史记录"""
    global history
    if len(history) >= max_history:
        history.pop(0)  # 移除最旧的记录
    history.append(img.copy())


def undo():
    """撤回上一步操作"""
    global img, history
    if len(history) > 0:
        img = history.pop()
        return True
    return False


def mouse_callback(event, x, y, flags, param):
    global drawing, last_point, mouse_x, mouse_y, scale, offset_x, offset_y

    # 转换为原始图像坐标
    orig_x = int((x + offset_x) / scale)
    orig_y = int((y + offset_y) / scale)
    mouse_x, mouse_y = orig_x, orig_y

    if event == cv2.EVENT_LBUTTONDOWN:
        drawing = True
        last_point = (orig_x, orig_y)
        save_to_history(img)  # 开始绘制时保存状态

    elif event == cv2.EVENT_MOUSEMOVE:
        if drawing and last_point is not None:
            thickness = max(1, int(2 * circle_radius))
            cv2.line(img, last_point, (orig_x, orig_y), (0, 0, 0), thickness)
            last_point = (orig_x, orig_y)

    elif event == cv2.EVENT_LBUTTONUP:
        drawing = False
        last_point = None

    elif event == cv2.EVENT_MOUSEWHEEL:
        # 获取鼠标在原始图像上的位置
        mouse_img_x = (x + offset_x) / scale
        mouse_img_y = (y + offset_y) / scale

        # 计算缩放前后的比例关系
        old_scale = scale
        if flags > 0:  # 放大
            scale += scale_step
        else:  # 缩小
            scale = max(0.1, scale - scale_step)

        # 调整偏移量，使鼠标位置保持相对不变
        offset_x = mouse_img_x * scale - x
        offset_y = mouse_img_y * scale - y

        # 限制偏移范围，防止图像移出窗口
        img_h, img_w = img.shape[:2]
        max_offset_x = max(0, img_w * scale - window_width)
        max_offset_y = max(0, img_h * scale - window_height)

        offset_x = np.clip(offset_x, 0, max_offset_x)
        offset_y = np.clip(offset_y, 0, max_offset_y)


cv2.namedWindow("Paint Window", cv2.WINDOW_NORMAL)
cv2.resizeWindow("Paint Window", window_width, window_height)
cv2.setMouseCallback("Paint Window", mouse_callback)

# 遍历图片文件夹
for img_name in os.listdir(image_dir):
    if not img_name.lower().endswith(('.jpg', '.png', '.jpeg')):
        continue

    img_path = os.path.join(image_dir, img_name)
    img = cv2.imread(img_path)
    if img is None:
        print(f"无法读取图片: {img_name}")
        continue

    # 重置参数
    scale = 1.0
    offset_x, offset_y = 0, 0
    history = []  # 清空历史记录
    save_to_history(img)  # 保存初始状态

    while True:
        # 创建临时图像副本
        temp_img = img.copy()

        # 绘制预览圆
        if 0 <= mouse_x < temp_img.shape[1] and 0 <= mouse_y < temp_img.shape[0]:
            cv2.circle(temp_img, (mouse_x, mouse_y), circle_radius, (0, 255, 0), 1)

        # 缩放图像
        if scale != 1.0:
            h, w = temp_img.shape[:2]
            temp_img = cv2.resize(temp_img, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_LINEAR)

        # 应用偏移并裁剪显示区域
        h, w = temp_img.shape[:2]
        x1 = int(offset_x)
        y1 = int(offset_y)
        x2 = min(x1 + window_width, w)
        y2 = min(y1 + window_height, h)

        if x1 < w and y1 < h:
            display_img = temp_img[y1:y2, x1:x2]
            # 如果裁剪后的图像小于窗口大小，填充黑色背景
            if display_img.shape[0] < window_height or display_img.shape[1] < window_width:
                border = np.zeros((window_height, window_width, 3), dtype=np.uint8)
                border[0:display_img.shape[0], 0:display_img.shape[1]] = display_img
                display_img = border
        else:
            display_img = np.zeros((window_height, window_width, 3), dtype=np.uint8)

        cv2.imshow("Paint Window", display_img)

        key = cv2.waitKey(1) & 0xFF
        if key == ord('+') or key == ord('='):
            circle_radius += 1
        elif key == ord('-') and circle_radius > 1:
            circle_radius -= 1
        elif key == 27:  # ESC
            exit(0)
        elif key == ord('s'):  # 保存
            output_img_path = os.path.join(output_dir, img_name)
            cv2.imwrite(output_img_path, img)
            print(f"已保存: {output_img_path}")
        elif key == ord('q'):  # 退出当前图片
            break
        elif key == ord('z'):  # Ctrl+Z 撤回
            if undo():
                print("已撤回上一步操作")
            else:
                print("无法撤回，历史记录为空")

cv2.destroyAllWindows()
print("所有图像处理完成！")